#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import pymysql
import pandas as pd
from sshtunnel import SSHTunnelForwarder
from config import SSH_HOST, SSH_PORT, SSH_USERNAME, SSH_PASSWORD, DB_HOST, DB_PORT, DB_USERNAME, DB_PASSWORD


def create_ssh_tunnel():
    """
    使用sshtunnel库创建SSH隧道，返回SSHTunnelForwarder对象和本地端口
    """
    try:
        tunnel = SSHTunnelForwarder(
            (SSH_HOST, SSH_PORT),
            ssh_username=SSH_USERNAME,
            ssh_password=SSH_PASSWORD,
            remote_bind_address=(DB_HOST, DB_PORT),
            local_bind_address=('127.0.0.1', 0)  # 自动分配本地端口
        )
        tunnel.start()
        local_port = tunnel.local_bind_port
        print(f"SSH隧道已建立，本地端口: {local_port}")
        return tunnel, local_port
    except Exception as e:
        print(f"创建SSH隧道失败: {e}")
        return None, None


def get_sales_data():
    """
    执行SQL查询获取小渠道和大渠道的营业额数据
    """
    tunnel, local_port = create_ssh_tunnel()
    if tunnel is None:
        print("Failed to create SSH tunnel. Exiting.")
        return None, None

    database = 'tll_bi_dw'
    try:
        conn = pymysql.connect(
            host='127.0.0.1',
            port=local_port,
            user=DB_USERNAME,
            password=DB_PASSWORD,
            database=database
        )
        print("Database connection established.")
        
        # 小渠道查询
        small_channel_sql = """
        SELECT substr(pt, 1, 6) AS 月份, shop_id as 门店编号, SUM(amount) as 营业额 
        FROM ads_dbs_trade_shop_pay_channel_di 
        WHERE substr(pt, 1, 4) >= '2024'
        GROUP BY shop_id,substr(pt, 1, 6)
        ORDER BY substr(pt, 1, 6)
        """
        
        # 大渠道查询
        large_channel_sql = """
        SELECT stat_shop_id as 门店编号, substr(business_date, 1, 6) AS 月份, SUM(total_amount) as 大渠道营业额 
        FROM ads_dbs_trade_shop_di 
        WHERE substr(business_date, 1, 4) >= '2024'
        GROUP BY stat_shop_id, substr(business_date, 1, 6) 
        ORDER BY substr(business_date, 1, 6)
        """
        
        # 执行查询
        print("正在执行小渠道查询...")
        small_channel_df = pd.read_sql(small_channel_sql, conn)
        print(f"小渠道查询完成，共{len(small_channel_df)}条记录")
        
        print("正在执行大渠道查询...")
        large_channel_df = pd.read_sql(large_channel_sql, conn)
        print(f"大渠道查询完成，共{len(large_channel_df)}条记录")
        
        return small_channel_df, large_channel_df
        
    except Exception as e:
        print(f"查询数据失败: {e}")
        return None, None
    finally:
        conn.close()
        tunnel.stop()
        print("数据库连接和SSH隧道已关闭")


def get_meituan_data():
    """
    读取美团营业额数据
    """
    try:
        # 读取Excel文件
        df = pd.read_excel('美团营业额.xlsx')
        
        # 处理门店编号，将类似"TLL05913-2"的格式统一为"TLL05913"
        df['机构编号'] = df['机构编号'].str.split('-').str[0]
        
        # 重命名列
        df.rename(columns={'营业月份': '月份', '机构编号': '门店编号', '营业额(元)': '美团营业额'}, inplace=True)
        
        # 格式化月份字段为YYYYMM格式
        df['月份'] = pd.to_datetime(df['月份']).dt.strftime('%Y%m')
        
        # 按门店编号和月份汇总数据
        meituan_df = df.groupby(['门店编号', '月份'])['美团营业额'].sum().reset_index()
        
        print(f"美团数据读取完成，共{len(meituan_df)}条记录")
        return meituan_df
        
    except Exception as e:
        print(f"读取美团数据失败: {e}")
        return None


def generate_report(small_channel_df, large_channel_df, meituan_df=None):
    """
    生成营业额差异报告并保存到Excel文件
    """
    if small_channel_df is None or large_channel_df is None:
        print("数据获取失败，无法生成报告")
        return
    
    # 合并数据
    # 重命名列以便合并
    small_channel_df.rename(columns={'门店编号': 'shop_id', '月份': 'month', '营业额': 'small_channel_sales'}, inplace=True)
    large_channel_df.rename(columns={'门店编号': 'shop_id', '月份': 'month', '大渠道营业额': 'large_channel_sales'}, inplace=True)
    
    # 如果有美团数据，也进行重命名
    if meituan_df is not None:
        meituan_df.rename(columns={'门店编号': 'shop_id', '月份': 'month', '美团营业额': 'meituan_sales'}, inplace=True)
    
    # 合并数据框
    merged_df = pd.merge(small_channel_df, large_channel_df, on=['shop_id', 'month'], how='outer')
    
    # 如果有美团数据，也合并进来
    if meituan_df is not None:
        merged_df = pd.merge(merged_df, meituan_df, on=['shop_id', 'month'], how='outer')
    
    # 计算差异
    merged_df['营业额差异'] = merged_df['large_channel_sales'] - merged_df['small_channel_sales']
    
    # 重新排列列的顺序
    if meituan_df is not None:
        final_df = merged_df[['month', 'shop_id', 'small_channel_sales', 'large_channel_sales', 'meituan_sales', '营业额差异']]
    else:
        final_df = merged_df[['month', 'shop_id', 'small_channel_sales', 'large_channel_sales', '营业额差异']]
    
    # 保存到Excel文件
    output_file = 'sales_report.xlsx'
    with pd.ExcelWriter(output_file, engine='openpyxl') as writer:
        final_df.to_excel(writer, sheet_name='营业额对比', index=False)
        small_channel_df.to_excel(writer, sheet_name='小渠道详情', index=False)
        large_channel_df.to_excel(writer, sheet_name='大渠道详情', index=False)
        if meituan_df is not None:
            meituan_df.to_excel(writer, sheet_name='美团详情', index=False)
    
    print(f"报告已生成并保存到 {output_file}")


def main():
    """
    主函数
    """
    small_channel_df, large_channel_df = get_sales_data()
    meituan_df = get_meituan_data()
    generate_report(small_channel_df, large_channel_df, meituan_df)


if __name__ == "__main__":
    main()